import pygame as pg
import pymunk
import pymunk.batch
import random
from math import degrees # 用于将弧度转换为角度

def create_wall(x, y, w, h):
    '''创建墙壁'''
    body = pymunk.Body(body_type=pymunk.Body.STATIC) # 墙壁是静态的身体
    body.position = x, y # 设置位置
    shape = pymunk.Poly(body, [(0, 0), (w, 0), (w, h), (0, h)]) # 多边形
    shape.elasticity = 0.5 # 设置弹性
    shape.friction = 0.7 # 设置摩擦系数
    shape.collision_type = 99 # 设置碰撞类别
    space.add(body, shape) # 加入空间

def create_ball(value):
    '''创建小球并返回Surface和Rect对象'''
    ball_type = int(2 ** value)
    size = (value + 1) * 15 + 10 # 小球大小
    image = pg.Surface((size, size)).convert_alpha()
    image.fill((0, 0, 0, 0))
    rect = image.get_rect()
    pg.draw.circle(image, all_sizes[ball_type], (size/2, size/2), size/2)

    font = pg.font.Font("font.ttf", size//3) # 字体
    r = font.render(str(ball_type), True, (255, 255, 255)) # 文字渲染
    image.blit(r, r.get_rect(center=pg.Vector2(rect.size)/2))

    return image, rect
        
def ball_hit(arbiter, space, data):
    '''小球发生碰撞的post_solve'''
    a, b = arbiter.shapes # 发生碰撞的两个小球的形状
    if getattr(a, "value", "a") == getattr(b, "value", "b"): # 两个小球类别相同就进行合成
        if a.value + 1 < len(all_sizes): # 小球不能为最大的2048
            space.remove(a, a.body, b, b.body)
            a.sprite.kill()
            b.sprite.kill() # 删除原本的小球
            
            all_sprites.add(Ball(a.value + 1, *arbiter.contact_point_set.points[0].point_a, True)) # 在原来的位置添加新的更大的小球
            if a.value - max(current_values) > 1:
                current_values.append(max(current_values) + 1) # 可以提供更大的小球

class Sprite(pg.sprite.Sprite):
    def draw(self):
        screen.blit(self.image, self.rect)

class Ball(Sprite):
    def __init__(self, value, pos_x, pos_y, scale=False):
        global score
        
        super().__init__()
        score += list(all_sizes)[value] # 增加得分

        self.value = value
        self.image, self.rect = create_ball(value)
        self.orig_image = self.image.copy() # 便于缩放小球的图片
        size = self.image.get_width()

        self.body = pymunk.Body()
        self.body.position = (pg.math.clamp(pos_x, size//2, WIDTH - size//2), pos_y) # 把小球的位置限制在屏幕中
        self.radius = size // 2
        self.current_radius = 1 if scale else self.radius # scale表示是否让小球从小变大
        self.shape = pymunk.Circle(self.body, self.current_radius)
        self.shape.mass = value + 1
        self.shape.elasticity = 0.5
        self.shape.friction = 0.3
        space.add(self.body, self.shape) # 添加小球的身体和形状
        
        self.shape.sprite = self
        self.shape.value = self.value

    def update(self):
        self.rect.center = self.body.position

        if self.current_radius < self.radius:
            self.current_radius += 3 # 增加小球的半径
            if self.current_radius > self.radius:
                self.current_radius = self.radius
            self.shape.unsafe_set_radius(self.current_radius) # 重设小球的半径
            self.image = pg.transform.scale_by(self.orig_image, self.current_radius/self.radius) # 缩放小球的图片

    def draw(self):
        image = pg.transform.rotate(self.image, -degrees(self.body.angle)) # 根据body的角度旋转小球
        self.rect = image.get_rect(center=self.rect.center)
        screen.blit(image, self.rect)

class BallTip(Sprite):
    POS_Y = 10

    def __init__(self):
        super().__init__()
        self.update_value(0)

    @property
    def not_ready(self):
        return self.rect.y < self.POS_Y
        
    def update_value(self, value): # 更新添加的小球
        self.value = value
        self.image, self.rect = create_ball(value)
        self.rect.bottom = -55

    def update(self):
        self.rect.centerx = pg.mouse.get_pos()[0] # 将小球的中心x坐标设为鼠标x坐标

        if self.not_ready:
            self.rect.y += 2 # 实现小球从顶端进入屏幕的动画效果
            if self.rect.y > self.POS_Y:
                self.rect.y = self.POS_Y

class Score(Sprite):
    def __init__(self):
        super().__init__()
        self.value = 0
        self += 0

    def __add__(self, score): # Score对象支持加法运算
        self.value += score
        self.image = ui_font.render(f"Score: {self.value}", True, (0, 0, 0)) # 更新得分文字
        self.rect = self.image.get_rect(topleft=(5, 5))
        return self
    
WIDTH = 500
HEIGHT = 650 # 窗口尺寸
FPS = 60 # 帧速率
BG = (245, 245, 245) # 背景色

pg.init()
screen = pg.display.set_mode((WIDTH, HEIGHT)) # 设置窗口尺寸
pg.display.set_caption("2048") # 设置标题
clock = pg.time.Clock()
ui_font = pg.font.Font("font.ttf", 16)

all_sizes = {
    1:(88, 88, 88),
    2:(255, 128, 0),
    4:(128, 255, 255),
    8:(128, 0, 255),
    16:(128, 128, 0),
    32:(0, 255, 0),
    64:(255, 128, 255),
    128:(255, 255, 0),
    256:(255, 128, 100),
    512:(64, 128, 128),
    1024:(190, 120, 120),
    2048:(0, 0, 0)
} # 不同类型小球的颜色

while True:
    all_sprites = pg.sprite.Group()
    current_values = [0] # 当前可以用于添加的小球类别

    ball_tip = BallTip()
    all_sprites.add(ball_tip)
    score = Score()
    all_sprites.add(score)

    space = pymunk.Space() # 创建空间
    space.gravity = (0, 980) # 设置重力
    space.sleep_time_threshold = 1 # 1s空闲后身体进入睡眠状态

    create_wall(-10, 0, 10, HEIGHT)
    create_wall(WIDTH, 0, 10, HEIGHT)
    create_wall(0, HEIGHT, WIDTH, 10) # 创建位于左、右、下侧的墙壁

    handler = space.add_collision_handler(0, 0) # 小球和小球发生碰撞（碰撞类别默认为0）
    handler.post_solve = ball_hit # 发生碰撞时调用ball_hit进行处理

    buffer = pymunk.batch.Buffer() # 用于储存小球位置的缓存

    running = True
    while running: # 游戏
        screen.fill(BG)

        all_sprites.update()

        buffer.clear() # 清空缓存
        pymunk.batch.get_space_bodies(
            space,
            pymunk.batch.BodyFields.POSITION,
            buffer,
        ) # 获得空间中所有身体的position
        data = tuple(memoryview(buffer.float_buf()).cast("d")) # 批处理所有小球的位置信息
        for i in range(1, len(data), 2):
            if data[i] < 0: # 小球的中心点y坐标小于0
                running = False # 游戏失败
                break

        for sprite in all_sprites:
            sprite.draw() # 绘制所有精灵
        
        for event in pg.event.get(): 
            if event.type == pg.QUIT: # 点击关闭按钮
                pg.quit()
                raise SystemExit
            elif event.type == pg.MOUSEBUTTONDOWN and event.button == 1: # 点击鼠标左键
                if not ball_tip.not_ready:
                    all_sprites.add(Ball(ball_tip.value, event.pos[0], ball_tip.rect.centery))
                    ball_tip.update_value(random.choice(current_values)) # 从可以添加的小球里面随机选择一个

        space.step(1 / FPS) # 运行空间
        clock.tick(FPS) # 控制帧率
        pg.display.flip() # 刷新pygame窗口

    score = score.value
    tip = ui_font.render(f"Game Over!\nScore: {score}\n\nPress R to retry", True, (0, 0, 0))
    tip_rect = tip.get_rect(center=(WIDTH//2, HEIGHT//2)) # 游戏结束提示

    running = True
    while running: # 游戏结束画面
        screen.fill(BG)
        screen.blit(tip, tip_rect)
        
        for event in pg.event.get(): 
            if event.type == pg.QUIT:
                pg.quit()
                raise SystemExit # 退出程序
            elif event.type == pg.KEYDOWN and event.key == pg.K_r:
                running = False # 按R键重试

        clock.tick(FPS)
        pg.display.flip()

